#!/usr/bin/env python3
import rospy
import cv2
import numpy as np
import threading
import time
from queue import Queue, Empty
from concurrent.futures import ThreadPoolExecutor
from cv_bridge import CvBridge
from sensor_msgs.msg import Image
from collections import deque
import os
import sys

# 添加YoloSegV3路径
sys.path.append(os.path.join(os.path.dirname(__file__), 'YoloSegV3'))

from rknn_yolo_executor import RKNNYoloExecutor
from yolo_processor import YoloProcessor
from visualizer import Visualizer

class RealTimeImageProcessor:
    def __init__(self):
        rospy.init_node('realtime_image_processor', anonymous=True)
        
        # 初始化参数
        self.bridge = CvBridge()
        self.max_queue_size = 10  # 最大队列长度
        self.process_interval = 0.1  # 处理间隔（秒）
        
        # 创建队列
        self.raw_image_queue = Queue(maxsize=self.max_queue_size)  # 原始图像队列
        self.detection_queue = Queue(maxsize=self.max_queue_size)  # 检测结果队列
        self.result_queue = Queue(maxsize=self.max_queue_size)     # 最终结果队列
        
        # 模型路径
        self.pose_model_path = os.path.join(os.path.dirname(__file__), 'models/merge_yolov8_pose.rknn')
        self.seg_model_path = os.path.join(os.path.dirname(__file__), 'models/merge_yolov8_seg.rknn')
        
        # 初始化YOLO执行器
        rospy.loginfo("初始化YOLO模型...")
        self.yolo_executor = RKNNYoloExecutor(
            pose_model_path=self.pose_model_path,
            seg_model_path=self.seg_model_path,
            pose_npu_ids=[0],  # Pose模型使用NPU Core 0
            seg_npu_ids=[1]    # Seg模型使用NPU Core 1
        )
        
        # 初始化处理器和可视化器
        self.processor = YoloProcessor()
        self.visualizer = Visualizer()
        
        # 设置标准位置（可以从参数获取或使用默认值）
        self.setup_standard_position()
        
        # 线程控制
        self.running = True
        self.threads = []
        
        # 订阅ROS话题
        rospy.Subscriber('/camera/color/image_raw', Image, self.color_callback, queue_size=1)
        rospy.Subscriber('/camera/depth/image_raw', Image, self.depth_callback, queue_size=1)
        
        # 存储最新的深度图
        self.latest_depth = None
        self.depth_lock = threading.Lock()
        
        # 启动处理线程
        self.start_threads()
        
        rospy.loginfo("实时图像处理系统启动完成")

    def setup_standard_position(self):
        """设置标准位置 - 可以从参数文件读取或使用默认值"""
        standard_vector = rospy.get_param('~standard_vector', None)
        if standard_vector:
            self.processor.set_standard_vector(standard_vector)
            rospy.loginfo("从参数加载标准位置")
        else:
            rospy.logwarn("未设置标准位置，将使用第一帧作为标准位置")

    def color_callback(self, msg):
        """彩色图像回调函数"""
        try:
            cv_image = self.bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
            
            # 获取对应的深度图
            with self.depth_lock:
                depth_image = self.latest_depth.copy() if self.latest_depth is not None else None
            
            if depth_image is not None:
                # 将图像对放入队列
                image_data = {
                    'color': cv_image,
                    'depth': depth_image,
                    'timestamp': msg.header.stamp
                }
                
                try:
                    self.raw_image_queue.put_nowait(image_data)
                except:
                    # 如果队列满了，丢弃最旧的数据
                    try:
                        self.raw_image_queue.get_nowait()
                        self.raw_image_queue.put_nowait(image_data)
                    except:
                        pass
        except Exception as e:
            rospy.logerr(f"彩色图像处理错误: {e}")

    def depth_callback(self, msg):
        """深度图像回调函数"""
        try:
            cv_image = self.bridge.imgmsg_to_cv2(msg, desired_encoding='passthrough')
            with self.depth_lock:
                self.latest_depth = cv_image
        except Exception as e:
            rospy.logerr(f"深度图像处理错误: {e}")

    def detection_thread(self):
        """检测线程 - 运行YOLO模型"""
        rospy.loginfo("检测线程启动")
        
        while self.running and not rospy.is_shutdown():
            try:
                # 从原始图像队列获取数据
                image_data = self.raw_image_queue.get(timeout=1.0)
                
                color_image = image_data['color']
                
                # 提交给YOLO执行器
                self.yolo_executor.submit_pose(color_image)
                self.yolo_executor.submit_seg(color_image)
                
                # 获取检测结果
                pose_result = self.yolo_executor.get_pose_result()
                seg_result = self.yolo_executor.get_seg_result()
                
                if pose_result is not None and seg_result is not None:
                    detection_data = {
                        'color': color_image,
                        'depth': image_data['depth'],
                        'pose_result': pose_result,
                        'seg_result': seg_result,
                        'timestamp': image_data['timestamp']
                    }
                    
                    try:
                        self.detection_queue.put_nowait(detection_data)
                    except:
                        # 队列满了，丢弃最旧的数据
                        try:
                            self.detection_queue.get_nowait()
                            self.detection_queue.put_nowait(detection_data)
                        except:
                            pass
                
            except Empty:
                continue
            except Exception as e:
                rospy.logerr(f"检测线程错误: {e}")

    def processing_thread(self):
        """处理线程 - 计算偏移和旋转"""
        rospy.loginfo("处理线程启动")
        
        while self.running and not rospy.is_shutdown():
            try:
                # 从检测结果队列获取数据
                detection_data = self.detection_queue.get(timeout=1.0)
                
                # 处理检测结果
                result = self.processor.process_detection(
                    color_image=detection_data['color'],
                    depth_image=detection_data['depth'],
                    pose_result=detection_data['pose_result'],
                    seg_result=detection_data['seg_result']
                )
                
                if result is not None:
                    result_data = {
                        'color': detection_data['color'],
                        'depth': detection_data['depth'],
                        'result': result,
                        'timestamp': detection_data['timestamp']
                    }
                    
                    try:
                        self.result_queue.put_nowait(result_data)
                    except:
                        # 队列满了，丢弃最旧的数据
                        try:
                            self.result_queue.get_nowait()
                            self.result_queue.put_nowait(result_data)
                        except:
                            pass
                
            except Empty:
                continue
            except Exception as e:
                rospy.logerr(f"处理线程错误: {e}")

    def visualization_thread(self):
        """可视化线程 - 显示结果"""
        rospy.loginfo("可视化线程启动")
        
        # 创建显示窗口
        cv2.namedWindow('Battery Detection Result', cv2.WINDOW_NORMAL)
        cv2.resizeWindow('Battery Detection Result', 1280, 720)
        
        while self.running and not rospy.is_shutdown():
            try:
                # 从结果队列获取数据
                result_data = self.result_queue.get(timeout=1.0)
                
                # 可视化结果
                display_image = self.visualizer.draw_result(
                    image=result_data['color'],
                    result=result_data['result']
                )
                
                # 显示图像
                cv2.imshow('Battery Detection Result', display_image)
                
                # 处理键盘事件
                key = cv2.waitKey(1) & 0xFF
                if key == ord('q'):
                    rospy.loginfo("用户请求退出")
                    self.running = False
                    break
                elif key == ord('s'):
                    # 保存当前标准位置
                    self.processor.save_current_as_standard(result_data['result'])
                    rospy.loginfo("保存当前位置为标准位置")
                
                # 打印结果信息
                if 'vector' in result_data['result']:
                    vector = result_data['result']['vector']
                    rospy.loginfo_throttle(1.0, f"偏差和夹角: {vector}")
                
            except Empty:
                # 显示空白图像或保持最后一帧
                cv2.waitKey(30)
                continue
            except Exception as e:
                rospy.logerr(f"可视化线程错误: {e}")
        
        cv2.destroyAllWindows()

    def start_threads(self):
        """启动所有处理线程"""
        # 检测线程
        detection_t = threading.Thread(target=self.detection_thread, daemon=True)
        detection_t.start()
        self.threads.append(detection_t)
        
        # 处理线程
        processing_t = threading.Thread(target=self.processing_thread, daemon=True)
        processing_t.start()
        self.threads.append(processing_t)
        
        # 可视化线程
        visualization_t = threading.Thread(target=self.visualization_thread, daemon=True)
        visualization_t.start()
        self.threads.append(visualization_t)

    def shutdown(self):
        """关闭系统"""
        rospy.loginfo("正在关闭实时图像处理系统...")
        self.running = False
        
        # 等待线程结束
        for thread in self.threads:
            thread.join(timeout=2.0)
        
        # 释放YOLO执行器
        self.yolo_executor.release()
        
        rospy.loginfo("系统关闭完成")

if __name__ == '__main__':
    processor = None
    try:
        processor = RealTimeImageProcessor()
        rospy.spin()
    except KeyboardInterrupt:
        rospy.loginfo("接收到键盘中断信号")
    except Exception as e:
        rospy.logerr(f"主程序错误: {e}")
    finally:
        if processor:
            processor.shutdown()